import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from umap import UMAP  # 直接导入UMAP类（修正导入方式）

# 配置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Heiti TC', 'STHeiti']  # 支持中文显示
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 加载数据集
iris = load_iris()
X, y = iris.data, iris.target
feature_names = iris.feature_names
target_names = iris.target_names

# 创建画布（2行3列布局）
plt.figure(figsize=(20, 12))

# --------------------------
# 原始数据可视化（前两个特征）
# --------------------------
plt.subplot(2, 3, 1)
plt.scatter(X[:, 0], X[:, 1], c=y, edgecolor='k', s=50)
plt.title('原始数据（前两个特征）')
plt.xlabel(feature_names[0])
plt.ylabel(feature_names[1])
plt.colorbar(ticks=[0, 1, 2], label='类别')

# --------------------------
# PCA不同维度对比
# --------------------------
for i, n in enumerate([1, 2, 3]):
    pca = PCA(n_components=n)
    X_pca = pca.fit_transform(X)

    # 绘制方差解释率
    if i == 0:
        plt.subplot(2, 3, 4)
        plt.bar(range(1, 4), pca.explained_variance_ratio_.cumsum(),
                alpha=0.5, color=['#1f77b4', '#ff7f0e', '#2ca02c'])
        plt.title('PCA方差解释率')
        plt.xticks(range(1, 4), ['1维', '2维', '3维'])
        plt.ylim(0, 1.1)

    # 绘制降维结果（当n=2时）
    if n == 2:
        plt.subplot(2, 3, 5)
        plt.scatter(X_pca[:, 0], X_pca[:, 1], c=y, edgecolor='k', s=50)
        plt.title('PCA 2D降维结果')
        plt.xlabel(f'主成分1 ({pca.explained_variance_ratio_[0]:.1%})')
        plt.ylabel(f'主成分2 ({pca.explained_variance_ratio_[1]:.1%})')
        plt.colorbar(ticks=[0, 1, 2], label='类别')

# --------------------------
# t-SNE降维
# --------------------------
tsne = TSNE(n_components=2, perplexity=30, random_state=24)
X_tsne = tsne.fit_transform(X)
plt.subplot(2, 3, 6)
plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y, edgecolor='k', s=50)
plt.title('t-SNE降维结果')
plt.xlabel('t-SNE维度1')
plt.ylabel('t-SNE维度2')
plt.colorbar(ticks=[0, 1, 2], label='类别')

# --------------------------
# UMAP降维（修正导入后）
# --------------------------
umap_model = UMAP(n_neighbors=15, min_dist=0.3, random_state=24)  # 直接使用UMAP类
X_umap = umap_model.fit_transform(X)
plt.subplot(2, 3, 2)
plt.scatter(X_umap[:, 0], X_umap[:, 1], c=y, edgecolor='k', s=50)
plt.title('UMAP降维结果')
plt.xlabel('UMAP维度1')
plt.ylabel('UMAP维度2')
plt.colorbar(ticks=[0, 1, 2], label='类别')

# --------------------------
# 多方法对比
# --------------------------
plt.subplot(2, 3, 3)
plt.scatter(X_pca[:, 0], X_pca[:, 1], c=y, edgecolor='k', s=50,
            label='PCA (2D)')
plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=y, edgecolor='r', s=50, marker='^',
            label='t-SNE')
plt.scatter(X_umap[:, 0], X_umap[:, 1], c=y, edgecolor='g', s=50, marker='s',
            label='UMAP')
plt.title('多方法对比')
plt.xlabel('第一维度')
plt.ylabel('第二维度')
plt.legend()

plt.tight_layout()
plt.show()